# -*- coding: utf-8 -*-
'''
Created on 15.11.2019

@author: yu03
'''
import datetime
import numpy as np
import os
import re
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import glob
from mpl_toolkits.mplot3d import Axes3D
from PyUeye_Unified.Full_Frame_Raw10 import folder_path, result_names, hor_index, hor_lines, ver_index, ver_lines
import sys
from FFT_Interpolation import FFT_interpolation_nonlinearity_compare
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from scipy import signal



# plt_show_mode = 'show' ### plt.show()
plt_show_mode = 'save' ### plt.save()

'''
    读取数据 (4通道)
'''
def Read_Data_4Ch(name):
    '''
        Return Data in File (4 Channels: Data_Ch1, Data_Ch2, Data_Ch3, Data_Ch4)
        File name required (default path)
    '''
    print('Reading Data')
    with open(name,'r') as fid:
        line=''
        while line[0:4] != '----':
            line = fid.readline()
            print(line)
            if line[0:2] == 'Fs':
                p, q, m, n = line.strip().split(' ')
                Fs = float(m)
                print('Fs = %f\n'%Fs)
        out_str = fid.readlines()
    Data_Ch1, Data_Ch2, Data_Ch3, Data_Ch4 = [], [], [], []
    for line in out_str:
        a, b, c, d= line.strip().split(', ')
        Data_Ch1.append(float(a))
        Data_Ch2.append(float(b))
        Data_Ch3.append(float(c))
        Data_Ch4.append(float(d))    
    Data_Ch1 = np.array(Data_Ch1)
    Data_Ch2 = np.array(Data_Ch2)
    Data_Ch3 = np.array(Data_Ch3)
    Data_Ch4 = np.array(Data_Ch4)
    return Data_Ch1, Data_Ch2, Data_Ch3, Data_Ch4, Fs

def Export_Data(out_name, header, out_str):
    print('Writing Data')
    with open(out_name,'w') as fid: ######################################################################################
        fid.writelines(header)
        fid.writelines(out_str)
    print('Finish Writing')

i = 0
fs_cam = 50
Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0
''' 
    for all groups
''' 
for result in result_names:
    ''' 
        reading .npy file
    ''' 
    file_name = result.split('\\')[-1] ### x_lines.py
#     file_name = re.findall(r"openmode\\(.+).npy", np_result)[0]
    print(file_name)
#     txt_name = folder_path + '\\' + file_name.split('.')[0] + '_results.txt'
#     PSD_name = folder_path + '\\' + file_name.split('.')[0] + '_PSD.txt'
    i += 1
    
    '''
        reading data
    '''
    
    now = datetime.datetime.now()
    hor_f_fit_set, hor_phase_centers, ver_f_fit_set, ver_phase_centers, fs = Read_Data_4Ch(result)
    
    hor_angle = (V_x-Lamda*hor_f_fit_set/2)*1e6 ### urad
    hor_length = np.unwrap(hor_phase_centers)/4/np.pi*Lamda*1e9 ### nm
    
    ver_angle = (V_y-Lamda*ver_f_fit_set/2)*1e6 ### urad
    ver_length = np.unwrap(ver_phase_centers)/4/np.pi*Lamda*1e9 ### nm
    
    hor_angle_avr = hor_angle - np.mean(hor_angle) 
    hor_length_avr = hor_length - np.mean(hor_length)
    ver_angle_avr = ver_angle - np.mean(ver_angle) 
    ver_length_avr = ver_length - np.mean(ver_length)
    
    f_line, PS_hor_angle_avr = signal.welch(hor_angle_avr, fs_cam, nperseg=len(hor_angle_avr))
    PSD_hor_angle_avr = np.sqrt(PS_hor_angle_avr)
    f_line, PS_hor_length_avr = signal.welch(hor_length_avr, fs_cam, nperseg=len(hor_length_avr))
    PSD_hor_length_avr = np.sqrt(PS_hor_length_avr)
    
    f_line, PS_ver_angle_avr = signal.welch(ver_angle_avr, fs_cam, nperseg=len(ver_angle_avr))
    PSD_ver_angle_avr = np.sqrt(PS_ver_angle_avr)
    f_line, PS_ver_length_avr = signal.welch(ver_length_avr, fs_cam, nperseg=len(ver_length_avr))
    PSD_ver_length_avr = np.sqrt(PS_ver_length_avr)
    
    ''' 
        Plotting averaged results
    '''   
    fig1 = plt.figure('Noise & PSD')
    plt.gcf().set_size_inches(18,9)
    
    ax1 = fig1.add_subplot(2, 2, 1)
    ax2 = fig1.add_subplot(2, 2, 2)
    ax3 = fig1.add_subplot(2, 2, 3)
    ax4 = fig1.add_subplot(2, 2, 4)
    
    ax1.plot(hor_angle_avr)
    ax2.plot(ver_angle_avr[:900])
    ax3.plot(hor_length_avr)
    ax4.plot(ver_length_avr)
    
    ax1.title.set_text('Hor. Tilt')
    ax2.title.set_text('Ver. Tilt')
    ax3.title.set_text('Hor. Displacement')
    ax4.title.set_text('Ver. Displacement')
    
    ax1.set_xlabel('Frame num.')
    ax2.set_xlabel('Frame num.')
    ax3.set_xlabel('Freq. \ Hz')
    ax4.set_xlabel('Freq. \ Hz')
    
    ax1.set_ylabel('Angle / urad')
    ax2.set_ylabel('Angle / urad')
    ax3.set_ylabel('Displacement / nm')
    ax4.set_ylabel('Displacement / nm')
    
    ax1.grid(which='both', axis='both')
    ax2.grid(which='both', axis='both')
    ax3.grid(which='both', axis='both')
    ax4.grid(which='both', axis='both')
    
    plt.tight_layout()
    if plt_show_mode == 'save':
        plt.savefig(folder_path+'\\'+file_name.split('.')[0]+'.png', dpi=300)
        plt.close()
    elif plt_show_mode == 'show':
        plt.show()
    else:
        sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)
    
#     ''' 
#         Results output
#     '''
#     if 1:
#         ''' 
#             Header
#         '''
#         header = ['Fs = %e (Hz)\n' %fs_cam,#########################################################################################################
#               'Channel_1: D\n',############################################################################################
#               'Channel_2: Angle hor.\n',############################################################################################
#               '-------------------------------------------------\n',
#               ]
#         ''' 
#             Exporting TXT
#         '''
#         out_str = ['%f, %f\n' %(hor_length_avr[i], hor_angle_avr[i]) for i in range(len(hor_angle_avr))]
#         Export_Data(txt_name, header, out_str)
#       
#         ''' 
#             PSD Header
#         '''
#         header_FFT = ['Fs = %e (Hz)\n' %50,#########################################################################################################
#               'Channel_1: PSD_f_line\n',############################################################################################
#               'Channel_2: PSD_D\n',############################################################################################
#               'Channel_3: PSD_hor.\n',############################################################################################
#               '-------------------------------------------------\n',
#               ]
#         ''' 
#             PSD Exporting TXT
#         '''
#         out_str_FFT = ['%f, %f, %f\n' %(f_line[i], PSD_hor_length_avr[i], PSD_hor_angle_avr[i]) for i in range(len(PSD_hor_length_avr))]
#         Export_Data(PSD_name, header_FFT, out_str_FFT)
    
#     ''' 
#         Plotting averaged results
#     '''   
#     fig1 = plt.figure('Tilt Noise & PSD'Noise_PSD    plt.gcf().set_size_inches(18,9)
#     
#     ax1 = fig1.add_subplot(2, 2, 1)
#     ax2 = fig1.add_subplot(2, 2, 2)
#     ax3 = fig1.add_subplot(2, 2, 3)
#     ax4 = fig1.add_subplot(2, 2, 4)
#     
#     ax1.plot(hor_angle_avr)
#     ax2.plot(ver_angle_avr)
#     ax3.loglog(f_line, PSD_hor_angle_avr)
#     ax4.loglog(f_line, PSD_ver_angle_avr)
#     
#     ax1.title.set_text('Hor. Tilt')
#     ax2.title.set_text('Ver. Tilt')
#     ax3.title.set_text('Hor. Tilt PSD')
#     ax4.title.set_text('Ver. Tilt PSD')
#     
#     ax1.set_xlabel('Frame num.')
#     ax2.set_xlabel('Frame num.')
#     ax3.set_xlabel('Freq. \ Hz')
#     ax4.set_xlabel('Freq. \ Hz')
#     
#     ax1.set_ylabel('Angle / urad')
#     ax2.set_ylabel('Angle / urad')
#     ax3.set_ylabel('PSD / urad/sqrt(Hz)')
#     ax4.set_ylabel('PSD / urad/sqrt(Hz)')
#     
#     ax1.grid(which='both', axis='both')
#     ax2.grid(which='both', axis='both')
#     ax3.grid(which='both', axis='both')
#     ax4.grid(which='both', axis='both')
#     
#     plt.tight_layout()
#     if plt_show_mode == 'save':
#         plt.savefig(folder_path+'\\'+file_name.split('.')[0]+'_tilt.png', dpi=300)
#         plt.close()
#     elif plt_show_mode == 'show':
#         plt.show()
#     else:
#         sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)
#         
#         
#     fig1 = plt.figure('Displacement Noise & PSD'Noise_PSD    plt.gcf().set_size_inches(18,9)
#     
#     ax1 = fig1.add_subplot(2, 2, 1)
#     ax2 = fig1.add_subplot(2, 2, 2)
#     ax3 = fig1.add_subplot(2, 2, 3)
#     ax4 = fig1.add_subplot(2, 2, 4)
#     
#     ax1.plot(hor_length_avr)
#     ax2.plot(ver_length_avr)
#     ax3.loglog(f_line, PSD_hor_length_avr)
#     ax4.loglog(f_line, PSD_ver_length_avr)
#     
#     ax1.title.set_text('Hor. Displacement')
#     ax2.title.set_text('Ver. Displacement')
#     ax3.title.set_text('Hor. Displacement PSD')
#     ax4.title.set_text('Ver. Displacement PSD')
#     
#     ax1.set_xlabel('Frame num.')
#     ax2.set_xlabel('Frame num.')
#     ax3.set_xlabel('Freq. \ Hz')
#     ax4.set_xlabel('Freq. \ Hz')
#     
#     ax1.set_ylabel('Displacement / nm')
#     ax2.set_ylabel('Displacement / nm')
#     ax3.set_ylabel('PSD / nm/sqrt(Hz)')
#     ax4.set_ylabel('PSD / nm/sqrt(Hz)')
#     
#     ax1.grid(which='both', axis='both')
#     ax2.grid(which='both', axis='both')
#     ax3.grid(which='both', axis='both')
#     ax4.grid(which='both', axis='both')
#     
#     plt.tight_layout()
#     if plt_show_mode == 'save':
#         plt.savefig(folder_path+'\\'+file_name.split('.')[0]+'_displacement.png', dpi=300)
#         plt.close()
#     elif plt_show_mode == 'show':
#         plt.show()
#     else:
#         sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)
#     